#  -*- coding: utf-8 -*-

"""
普量学院量化投资课程系列案例源码包
普量学院版权所有
仅用于教学目的，严禁转发和用于盈利目的，违者必究
©Plouto-Quants All Rights Reserved

普量学院助教微信：niuxiaomi3
"""


from datetime import datetime, timedelta
import os
import time

import tushare as ts
from pymongo import UpdateOne

from database import DB_CONN
from stock_util import get_trading_dates, safe_float_conversion

"""
从tushare pro获取股票基础数据，保存到本地的MongoDB数据库中
升级版本：支持tushare pro API，优化执行效率
"""


def crawl_basic(begin_date=None, end_date=None, token=None):
    """
    抓取指定时间范围内的股票基础信息
    :param begin_date: 开始日期，格式：YYYY-MM-DD
    :param end_date: 结束日期，格式：YYYY-MM-DD
    :param token: tushare pro的token
    """

    # 如果没有指定开始日期，则默认为前一日
    if begin_date is None:
        begin_date = (datetime.now() - timedelta(days=1)).strftime('%Y-%m-%d')

    # 如果没有指定结束日期，则默认为前一日
    if end_date is None:
        end_date = (datetime.now() - timedelta(days=1)).strftime('%Y-%m-%d')

    # 获取指定日期范围的所有交易日列表
    all_dates = get_trading_dates(begin_date, end_date)
    
    print(f"开始抓取股票基础信息，日期范围：{begin_date} 到 {end_date}")
    print(f"共计 {len(all_dates)} 个交易日")

    total_start_time = time.time()
    total_api_time = 0
    total_process_time = 0
    success_count = 0
    error_count = 0

    # 按照每个交易日抓取
    for i, date in enumerate(all_dates):
        try:
            # 抓取当日的基本信息并返回时间统计
            api_time, process_time = crawl_basic_at_date(date, token=token)
            total_api_time += api_time
            total_process_time += process_time
            success_count += 1
            
            # 每10个交易日输出一次进度
            if (i + 1) % 10 == 0:
                print(f"进度：{i+1}/{len(all_dates)}，成功：{success_count}，失败：{error_count}")
                
        except Exception as e:
            print('抓取股票基本信息时出错，日期：%s，错误：%s' % (date, str(e)), flush=True)
            error_count += 1

    total_time = time.time() - total_start_time
    
    # 输出性能统计
    print(f"\n=== 性能统计 ===")
    print(f"总耗时：{total_time:.2f}秒")
    print(f"API调用总耗时：{total_api_time:.2f}秒 ({total_api_time/total_time*100:.1f}%)")
    print(f"数据处理总耗时：{total_process_time:.2f}秒 ({total_process_time/total_time*100:.1f}%)")
    print(f"平均每个交易日：{total_time/len(all_dates):.2f}秒")
    print(f"平均API调用时间：{total_api_time/len(all_dates):.2f}秒")
    print(f"成功处理：{success_count}个交易日")
    print(f"失败：{error_count}个交易日")


def crawl_basic_at_date(date, token=None):
    """
    从Tushare pro抓取指定日期的股票基本信息
    :param date: 日期，格式：YYYY-MM-DD
    :param token: tushare pro的token
    :return: (api_time, process_time) 返回API调用时间和数据处理时间
    """

    if token is None:
        token = os.getenv('TUSHARE_TOKEN')

    if not token:
        raise ValueError("请设置tushare pro token! 访问 https://tushare.pro/ 获取token")

    ts.set_token(token)
    pro = ts.pro_api()

    # 转换日期格式从YYYY-MM-DD到YYYYMMDD
    trade_date = date.replace('-', '')

    # 统计API调用时间
    api_start_time = time.time()
    
    # 使用tushare pro获取股票基本信息
    df_basics = pro.daily_basic(trade_date=trade_date)
    
    api_time = time.time() - api_start_time

    if df_basics is None or df_basics.empty:
        print(f'日期 {date} 没有基础信息数据')
        return api_time, 0

    # 统计数据处理时间
    process_start_time = time.time()
    
    # 处理tushare pro数据
    _process_pro_basic_data(df_basics, date)
    
    process_time = time.time() - process_start_time
    
    print(f'日期 {date}: API调用 {api_time:.2f}s, 数据处理 {process_time:.2f}s, 数据量 {len(df_basics)}条')
    
    return api_time, process_time


def _process_pro_basic_data(df_basics, date):
    """
    处理tushare pro的基本信息数据 - 优化版本
    """
    # 预分配列表大小以提高性能
    update_requests = []
    
    # 批量处理数据，减少逐行处理的开销
    for _, row in df_basics.iterrows():
        try:
            # 提取股票代码（去掉后缀）- 优化字符串操作
            ts_code = row['ts_code']
            code = ts_code[:6]  # 直接截取前6位，比split更快
            
            # 预构建文档，减少重复调用
            doc = {
                'code': code,
                'date': date,
                'total_mv': safe_float_conversion(row.get('total_mv'), 0),
                'circ_mv': safe_float_conversion(row.get('circ_mv'), 0),
                'pe': safe_float_conversion(row.get('pe'), 0),
                'pe_ttm': safe_float_conversion(row.get('pe_ttm'), 0),
                'pb': safe_float_conversion(row.get('pb'), 0),
                'ps': safe_float_conversion(row.get('ps'), 0),
                'ps_ttm': safe_float_conversion(row.get('ps_ttm'), 0),
                'total_share': safe_float_conversion(row.get('total_share'), 0),
                'float_share': safe_float_conversion(row.get('float_share'), 0),
                'free_share': safe_float_conversion(row.get('free_share'), 0),
                'turnover_rate': safe_float_conversion(row.get('turnover_rate'), 0),
                'turnover_rate_f': safe_float_conversion(row.get('turnover_rate_f'), 0),
                'volume_ratio': safe_float_conversion(row.get('volume_ratio'), 0),
            }

            # 直接添加到列表，避免重复创建UpdateOne对象
            update_requests.append(
                UpdateOne(
                    {'code': code, 'date': date},
                    {'$set': doc}, 
                    upsert=True
                )
            )

        except Exception as e:
            print('处理tushare pro数据时发生异常，股票代码：%s，日期：%s，错误：%s' % (ts_code, date, str(e)), flush=True)

    # 批量写入数据库
    if update_requests:
        try:
            # 使用ordered=False提高批量写入性能
            update_result = DB_CONN['basic'].bulk_write(update_requests, ordered=False)
            print('抓取股票基本信息，日期：%s, 插入：%4d条，更新：%4d条' %
                  (date, update_result.upserted_count, update_result.modified_count), flush=True)
        except Exception as e:
            print(f'批量写入数据库时出错，日期：{date}，错误：{str(e)}', flush=True)


# 新增：批量抓取优化版本
def crawl_basic_batch(begin_date=None, end_date=None, token=None, batch_size=5):
    """
    批量抓取股票基础信息 - 优化版本
    :param begin_date: 开始日期，格式：YYYY-MM-DD
    :param end_date: 结束日期，格式：YYYY-MM-DD
    :param token: tushare pro的token
    :param batch_size: 批量大小，每批处理多少个交易日
    """
    
    if begin_date is None:
        begin_date = (datetime.now() - timedelta(days=1)).strftime('%Y-%m-%d')

    if end_date is None:
        end_date = (datetime.now() - timedelta(days=1)).strftime('%Y-%m-%d')

    # 获取指定日期范围的所有交易日列表
    all_dates = get_trading_dates(begin_date, end_date)
    
    print(f"开始批量抓取股票基础信息，日期范围：{begin_date} 到 {end_date}")
    print(f"共计 {len(all_dates)} 个交易日，批量大小：{batch_size}")

    total_start_time = time.time()
    success_count = 0
    error_count = 0

    # 按批处理
    for i in range(0, len(all_dates), batch_size):
        batch_dates = all_dates[i:i + batch_size]
        
        try:
            # 批量处理这一批日期
            _process_date_batch(batch_dates, token)
            success_count += len(batch_dates)
            
            print(f"批量处理完成：{i+1}-{min(i+batch_size, len(all_dates))}/{len(all_dates)}")
            
            # 避免请求过于频繁
            time.sleep(0.1)
            
        except Exception as e:
            print(f'批量处理时出错，日期批次：{batch_dates}，错误：{str(e)}', flush=True)
            error_count += len(batch_dates)

    total_time = time.time() - total_start_time
    
    print(f"\n=== 批量处理性能统计 ===")
    print(f"总耗时：{total_time:.2f}秒")
    print(f"平均每个交易日：{total_time/len(all_dates):.2f}秒")
    print(f"成功处理：{success_count}个交易日")
    print(f"失败：{error_count}个交易日")


def _process_date_batch(dates, token):
    """
    批量处理多个日期的数据
    """
    if token is None:
        token = os.getenv('TUSHARE_TOKEN')

    if not token:
        raise ValueError("请设置tushare pro token!")

    ts.set_token(token)
    pro = ts.pro_api()

    all_update_requests = []
    
    for date in dates:
        try:
            trade_date = date.replace('-', '')
            
            # API调用
            df_basics = pro.daily_basic(trade_date=trade_date)
            
            if df_basics is None or df_basics.empty:
                continue
                
            # 处理数据
            for _, row in df_basics.iterrows():
                try:
                    ts_code = row['ts_code']
                    code = ts_code[:6]
                    
                    doc = {
                        'code': code,
                        'date': date,
                        'total_mv': safe_float_conversion(row.get('total_mv'), 0),
                        'circ_mv': safe_float_conversion(row.get('circ_mv'), 0),
                        'pe': safe_float_conversion(row.get('pe'), 0),
                        'pe_ttm': safe_float_conversion(row.get('pe_ttm'), 0),
                        'pb': safe_float_conversion(row.get('pb'), 0),
                        'ps': safe_float_conversion(row.get('ps'), 0),
                        'ps_ttm': safe_float_conversion(row.get('ps_ttm'), 0),
                        'total_share': safe_float_conversion(row.get('total_share'), 0),
                        'float_share': safe_float_conversion(row.get('float_share'), 0),
                        'free_share': safe_float_conversion(row.get('free_share'), 0),
                        'turnover_rate': safe_float_conversion(row.get('turnover_rate'), 0),
                        'turnover_rate_f': safe_float_conversion(row.get('turnover_rate_f'), 0),
                        'volume_ratio': safe_float_conversion(row.get('volume_ratio'), 0),
                    }

                    all_update_requests.append(
                        UpdateOne(
                            {'code': code, 'date': date},
                            {'$set': doc}, 
                            upsert=True
                        )
                    )
                except Exception as e:
                    print(f'处理数据时出错: {e}')
                    continue
                    
        except Exception as e:
            print(f'处理日期 {date} 时出错: {e}')
            continue
    
    # 批量写入所有数据
    if all_update_requests:
        try:
            update_result = DB_CONN['basic'].bulk_write(all_update_requests, ordered=False)
            print(f'批量写入完成，插入：{update_result.upserted_count}条，更新：{update_result.modified_count}条')
        except Exception as e:
            print(f'批量写入数据库时出错：{str(e)}')


if __name__ == '__main__':
    # 使用tushare pro token
    TOKEN = "b8860eddcd1c2b8d4c2c31075fa172bb031f0d217755d346bf223fa8"

    print("开始抓取股票基本信息...")
    
    # 选择使用哪种方式：
    # 方式1：原始方式，带详细时间统计
    #crawl_basic('2015-10-08', '2025-07-05', token=TOKEN)
    crawl_basic('2025-01-01', '2025-07-05', token=TOKEN)

    # 方式2：批量优化版本（可选）
    # crawl_basic_batch('2015-10-08', '2025-07-05', token=TOKEN, batch_size=5)

    print("股票基本信息抓取完成!")
